from tensorflow.keras.callbacks import TensorBoard,Callback
import os
import datetime
from .tool import confusion_matrix

class confusion(Callback):
    def __init__(self,train_summary,val_summary,class_name):
        super(confusion, self).__init__()
        self.train_summary = train_summary
        self.val_summary = val_summary
        self.class_name = class_name

    def on_epoch_end(self, epoch,logs=None):
        val_matrix = logs["val_confusion_matrix"]
        train_matrix = logs["confusion_matrix"]
        confusion_matrix(self.train_summary,train_matrix,self.class_name,epoch)
        confusion_matrix(self.val_summary,val_matrix,self.class_name,epoch)


class TensorBoard_eager_writer(TensorBoard):

    def __init__(self,
                 log_dir='out/logs',
                 experment_name="gradient_tape",
                 histogram_freq=0,
                 write_graph=True,
                 write_images=False,
                 update_freq='epoch',
                 profile_batch=2,
                 embeddings_freq=0,
                 embeddings_metadata=None,
                 **kwargs):
        super(TensorBoard_eager_writer, self).__init__(log_dir=log_dir, histogram_freq=histogram_freq,
                                                       write_graph=write_graph, write_images=write_images,
                                                       update_freq=update_freq, profile_batch=profile_batch,
                                                       embeddings_freq=embeddings_freq,
                                                       embeddings_metadata=embeddings_metadata, **kwargs)

        current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        train_log_dir = os.path.join(log_dir, experment_name, current_time, 'train')
        val_log_dir = os.path.join(log_dir, experment_name, current_time, 'validation')

        self._train_dir = train_log_dir
        self._val_dir = val_log_dir

        self._writers = {}
        self.train_writer = self._train_writer
        self.val_writer = self._val_writer

    def set_model(self, model):
        """Sets Keras model and writes graph if specified."""
        self.model = model

        self._train_step = self.model._train_counter
        self._val_step = self.model._test_counter

        if self.write_graph:
            self._write_keras_model_graph()
        if self.embeddings_freq:
            self._configure_embeddings()



